# coding=utf-8
# Copyright (C) xxx team - All Rights Reserved
#
# @Version:   3.9.2
# @Software:  PyCharm
# @FileName:  gen_dummy_data.py
# @CTime:     2021/5/3 16:56   
# @Author:    Haiyang Yu
# @Email:     yuys0602@163.com
# @UTime:     2021/5/3 16:56
#
# @Description:
#
#     xxx
#
import json
import os
import random
import codecs
import logging
from typing import List, Dict

logger = logging.getLogger(__name__)


def save_jsonl(data: List[Dict], file: str) -> None:
    """
    save data as json-ld format.

    Args:
        data: need saved data
        file: saved in file
    """
    with codecs.open(file, 'w', encoding='utf-8') as f:
        f.write(os.linesep.join(
            [json.dumps(line, ensure_ascii=False) for line in data]
        ))


def dummy_data(token_types: int = 3, num_samples: int = 100):
    """
    Args:
        token_types: token 类别总数
        num_samples: 样本数目，默认100条

    Returns:
       a list
    """
    lengths = [random.randint(5, 20) for _ in range(num_samples)]  # 句长

    data = []
    for lens in lengths:
        data.append({
            'ids': [101] + [random.randint(103, 200) for _ in range(lens)] + [102],
            'tags': [0] + [random.randint(0, token_types - 1) for _ in range(lens)] + [0],  # BIO 就3个
        })
    return data


def main(token_types: int = 3, num_samples: int = 1000, split_ratio: float = 0.8, need_fitting: bool = True):
    total_data = dummy_data(token_types, num_samples)
    split_main = int(split_ratio * num_samples)
    split_rest = num_samples - split_main
    train_data = total_data if need_fitting else total_data[:split_main]
    test_data = train_data[:split_rest] if need_fitting else total_data[split_main:]

    save_jsonl(train_data, 'train.txt')
    save_jsonl(test_data, 'test.txt')


if __name__ == '__main__':
    main(token_types=3, num_samples=1000, split_ratio=0.9, need_fitting=True)
